{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2020-11-16 11:22:24--  https://f000.backblazeb2.com/file/malaya-model/v38/generator/base.pb\n",
      "Resolving f000.backblazeb2.com (f000.backblazeb2.com)... 104.153.233.177\n",
      "Connecting to f000.backblazeb2.com (f000.backblazeb2.com)|104.153.233.177|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 \n",
      "Length: 1252668691 (1.2G) [application/octet-stream]\n",
      "Saving to: ‘base.pb’\n",
      "\n",
      "base.pb             100%[===================>]   1.17G  13.0MB/s    in 1m 46s  \n",
      "\n",
      "2020-11-16 11:24:12 (11.3 MB/s) - ‘base.pb’ saved [1252668691/1252668691]\n",
      "\n",
      "--2020-11-16 11:24:12--  https://f000.backblazeb2.com/file/malaya-model/v38/generator/small.pb\n",
      "Resolving f000.backblazeb2.com (f000.backblazeb2.com)... 104.153.233.177\n",
      "Connecting to f000.backblazeb2.com (f000.backblazeb2.com)|104.153.233.177|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 \n",
      "Length: 355570391 (339M) [application/octet-stream]\n",
      "Saving to: ‘small.pb’\n",
      "\n",
      "small.pb            100%[===================>] 339.10M  11.2MB/s    in 31s     \n",
      "\n",
      "2020-11-16 11:24:45 (11.1 MB/s) - ‘small.pb’ saved [355570391/355570391]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://f000.backblazeb2.com/file/malaya-model/v38/generator/base.pb\n",
    "!wget https://f000.backblazeb2.com/file/malaya-model/v38/generator/small.pb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.tools.graph_transforms import TransformGraph\n",
    "from glob import glob\n",
    "tf.set_random_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['small.pb', 'base.pb']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pbs = glob('*.pb')\n",
    "pbs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow_text\n",
    "import tf_sentencepiece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-6-ed0c3462ae0b>:12: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n",
      "small.pb\n",
      "base.pb\n"
     ]
    }
   ],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics, op=Dropout)',\n",
    "             'fold_constants(ignore_errors=true)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "#              'quantize_weights(fallback_min=-10, fallback_max=10)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']\n",
    "\n",
    "for pb in pbs:\n",
    "    input_graph_def = tf.GraphDef()\n",
    "    with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "        input_graph_def.ParseFromString(f.read())\n",
    "        \n",
    "    print(pb)\n",
    "    \n",
    "    transformed_graph_def = TransformGraph(input_graph_def, \n",
    "                                           ['inputs'],\n",
    "                                           ['SentenceTokenizer_1/SentenceTokenizer/SentencepieceDetokenizeOp'], transforms)\n",
    "    \n",
    "    with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "        f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_graph(frozen_graph_filename, **kwargs):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "\n",
    "    # https://github.com/onnx/tensorflow-onnx/issues/77#issuecomment-445066091\n",
    "    # to fix import T5\n",
    "    for node in graph_def.node:\n",
    "        if node.op == 'RefSwitch':\n",
    "            node.op = 'Switch'\n",
    "            for index in xrange(len(node.input)):\n",
    "                if 'moving_' in node.input[index]:\n",
    "                    node.input[index] = node.input[index] + '/read'\n",
    "        elif node.op == 'AssignSub':\n",
    "            node.op = 'Sub'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "        elif node.op == 'AssignAdd':\n",
    "            node.op = 'Add'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "        elif node.op == 'Assign':\n",
    "            node.op = 'Identity'\n",
    "            if 'use_locking' in node.attr:\n",
    "                del node.attr['use_locking']\n",
    "            if 'validate_shape' in node.attr:\n",
    "                del node.attr['validate_shape']\n",
    "            if len(node.input) == 2:\n",
    "                node.input[0] = node.input[1]\n",
    "                del node.input[1]\n",
    "\n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# g = load_graph('base.pb.quantized')\n",
    "# x = g.get_tensor_by_name('import/inputs:0')\n",
    "# logits = g.get_tensor_by_name('import/SentenceTokenizer_1/SentenceTokenizer/SentencepieceDetokenizeOp:0')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x, x_len, logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# test_sess.run(logits, feed_dict = {x: ['ringkasan: KUALA LUMPUR: Presiden Perancis Emmanuel Macron tidak menampakkan beliau seorang sosok yang bertamadun, selar Tun Dr Mahathir Mohamad menerusi kemas kini terbaharu di blognya. Bekas Perdana Menteri itu mendakwa, pemerintah tertinggi Perancis itu bersikap primitif kerana menuduh orang Islam terlibat dalam pembunuhan guru yang menghina Islam, malah menegaskan tindakan membunuh bukan ajaran Islam. Jelas Dr Mahathir, sejarah membuktikan bahawa orang Perancis pernah membunuh jutaan manusia, yang ramai mangsanya terdiri dari orang Islam.']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# test_sess.run(logits, feed_dict = {x: [[1,2,3,3,4]], x_len: [[1,1,1,1,1]]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['small.pb.quantized', 'base.pb.quantized']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "quantized = glob('*.pb.quantized')\n",
    "quantized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm *.pb*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(\n",
    "#     graph_def_file='test.pb',\n",
    "#     input_arrays=['Placeholder', 'Placeholder_1'],\n",
    "#     input_shapes={'Placeholder' : [None, 512], 'Placeholder_1': [None, 512]},\n",
    "#     output_arrays=['logits'],\n",
    "# )\n",
    "# # converter.allow_custom_ops=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.target_spec.supported_types = [tf.float16]\n",
    "# converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "# converter.experimental_new_converter = True\n",
    "# tflite_model = converter.convert()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, \n",
    "#                                        tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.target_spec.supported_types = [tf.float16]\n",
    "# converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "# tflite_model = converter.convert()\n",
    "\n",
    "# with open('tiny-bert-sentiment-float16.tflite', 'wb') as f:\n",
    "#     f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, \n",
    "#                                        tf.lite.OpsSet.SELECT_TF_OPS]\n",
    "# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]\n",
    "# tflite_model = converter.convert()\n",
    "\n",
    "# with open('tiny-bert-sentiment-hybrid.tflite', 'wb') as f:\n",
    "#     f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# interpreter = tf.lite.Interpreter(model_path='tiny-bert-sentiment-hybrid.tflite')\n",
    "# interpreter.allocate_tensors()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
